# -*- coding: utf-8 -*-
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# ===== 路径配置 =====
PATH = "/Users/linshangjin/25CCM/NKU-C/t2/aligned_quarterly_data2.csv"
OUT_CSV = "/Users/linshangjin/25CCM/NKU-C/t2/all_vars_yoy.csv"
OUT_PNG = "/Users/linshangjin/25CCM/NKU-C/t2/corr_heatmap_all_vars.png"

# 1. 读取数据
df = pd.read_csv(PATH, parse_dates=[0], index_col=0).sort_index()

# 2. 要计算同比增长率的列
cols_yoy_calc = ["CIER", "毕业生人数_万人", "GDP","第一产业占比", "第二产业占比", "第三产业占比"]

# 3. 对前三列计算季度同比增长率（百分比）
df_yoy_part = (df[cols_yoy_calc] - df[cols_yoy_calc].shift(4)) / df[cols_yoy_calc].shift(4) * 100

# 4. 其他列假定已是同比增长率
cols_keep = [c for c in df.columns if c not in cols_yoy_calc]
df_yoy_keep = df[cols_keep]

# 5. 合并
df_yoy = pd.concat([df_yoy_part, df_yoy_keep], axis=1)
df_yoy = df_yoy.dropna(how="any")  # 要求所有列都有值

# 6. 保存同比增长率表
df_yoy.to_csv(OUT_CSV, encoding="utf-8-sig", float_format="%.4f")

# 7. 计算 Pearson 相关系数
corr = df_yoy.corr(method="pearson")

# 8. 绘制热图（每格都有标注）
plt.figure(figsize=(8, 6))
sns.heatmap(corr, annot=True, fmt=".2f", cmap="RdBu_r", center=0,
            linewidths=0.5, cbar_kws={"shrink": .8}, annot_kws={"size": 10})
plt.title("Pearson Correlation Heatmap (YoY Growth Rates)", fontsize=14)
plt.tight_layout()
plt.savefig(OUT_PNG, dpi=150)
plt.close()

print(f"✅ 同比增长率表已保存: {OUT_CSV}")
print(f"✅ 相关性热图已保存: {OUT_PNG}")
